# Import necessary libraries
import json
import logging

# Import jieba for Chinese text segmentation
import jieba
# Import Rouge for evaluating text summarization
from rouge_chinese import Rouge

# Import dataset augmentation configurations
from myapp.configs.dataset_aug import DATASET_AUG_CONFIGS
# Import chat completions API
from myapp.third.maas.api import chat_completions
# Import exception logging utility
from myapp.utils.exception import log_exception

# Get a logger instance
log = logging.getLogger(__name__)
# Initialize a Rouge object
rouge = Rouge()


# Define a function to calculate ROUGE-L score
def calc_rouge_l(src_text, dst_text):
    """
    Calculate ROUGE-L score between two texts
    """
    # Segment the source text using jieba
    src_text = " ".join(jieba.cut(src_text))
    # Segment the destination text using jieba
    dst_text = " ".join(jieba.cut(dst_text))
    # Get ROUGE scores
    scores = rouge.get_scores(dst_text, src_text)
    # Return the ROUGE-L f-score
    return scores[0]['rouge-l']['f']


# Define a function to generate augmented data
def gen_augmented_data(origin_data):
    """
    Generate augmented data based on the given origin data
    """
    # Get top_p value from config, default to 0.9
    top_p = DATASET_AUG_CONFIGS.get("top_p", 0.9)
    # Get temperature value from config, default to 1.0
    temperature = DATASET_AUG_CONFIGS.get("temperature", 1.0)
    # Get ROUGE-L threshold from config, default to 0.8
    rouge_l_threshold = DATASET_AUG_CONFIGS.get("rouge_l_threshold", 0.8)
    # Get model code from config, default to "taichu_llm"
    model = DATASET_AUG_CONFIGS.get("model_code", "taichu_llm")
    # Get user template from config
    user_template = DATASET_AUG_CONFIGS.get("user_template", "###原始数据：\n{input}\n\n###增强数据:\n")

    # Call chat completions API to get augmented data
    res = chat_completions(user_template.format(input=origin_data)
                           , system_prompt=DATASET_AUG_CONFIGS.get("system_prompt"),
                           model=model, top_p=top_p, temperature=temperature)
    # Initialize a list to store the results
    ret = []
    # Initialize a list to store the best data
    best_data = []
    # Initialize the best score to a high value
    best_score = 10
    # Split the response into lines
    augmented_data = res.strip().split("\n")
    # Iterate over each augmented line
    for aug_line in augmented_data:
        # Remove the first 3 characters and strip whitespace
        augmented_string = aug_line.strip()[3:]
        # Log the augmented string
        log.info(f"augmented_string:{augmented_string}")
        # Use a try-except block to handle potential exceptions
        with log_exception:
            # Try to load the augmented string as JSON
            json.loads(augmented_string)
            # Calculate the ROUGE-L score
            score = calc_rouge_l(origin_data, augmented_string)
            # Log the score
            log.info(f"score:{score}")
            # Check if the current score is better than the best score
            if score < best_score:
                # Update the best score
                best_score = score
                # Update the best data
                best_data = [augmented_string]
            # Check if the score is below the threshold
            if score < rouge_l_threshold:
                # Append the augmented string to the result list
                ret.append(augmented_string)

    # If no data is below the threshold, return the best data
    if len(ret) == 0:
        return best_data

    # Return the list of augmented data
    return ret